import cv2
import os
import numpy as np
from tqdm import tqdm
import multiprocessing
from multiprocessing import Pool
from functools import partial

# 光流法参数
LK_PARAMS = dict(winSize=(15, 15),
                maxLevel=2,
                criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))

class ImageCropper:
    def __init__(self):
        self.crop_width = None  # 裁剪宽度
        self.reference_right_edge = None  # 参考右边缘位置
        self.edge_history = []  # 添加历史记录
        self.history_size = 5   # 历史帧数
        self.prev_gray = None  # 前一帧灰度图
        self.prev_points = None  # 前一帧特征点
        self.transform = np.eye(2, 3, dtype=np.float32)  # 变换矩阵
        
    def update_edge_history(self, new_edge):
        self.edge_history.append(new_edge)
        if len(self.edge_history) > self.history_size:
            self.edge_history.pop(0)
            
    def get_smoothed_edge(self):
        if not self.edge_history:
            return None
        return int(np.median(self.edge_history))  # 使用中位数平滑
            
    def detect_features(self, gray):
        # 使用Shi-Tomasi角点检测
        feature_params = dict(maxCorners=100,
                            qualityLevel=0.3,
                            minDistance=7,
                            blockSize=7)
        return cv2.goodFeaturesToTrack(gray, mask=None, **feature_params)
        
    def track_motion(self, curr_gray):
        if self.prev_gray is None:
            self.prev_gray = curr_gray
            self.prev_points = self.detect_features(curr_gray)
            return
            
        # 计算光流
        curr_points, status, _ = cv2.calcOpticalFlowPyrLK(
            self.prev_gray, curr_gray, self.prev_points, None, **LK_PARAMS)
            
        # 只保留好的特征点
        good_old = self.prev_points[status == 1]
        good_new = curr_points[status == 1]
        
        if len(good_old) >= 4:
            # 计算仿射变换矩阵
            self.transform, _ = cv2.estimateAffinePartial2D(
                good_old, good_new, method=cv2.RANSAC)
                
        # 更新前一帧信息
        self.prev_gray = curr_gray
        self.prev_points = good_new.reshape(-1, 1, 2)
        
    def find_right_edge(self, binary_img):
        # 计算垂直投影
        projection = np.sum(binary_img, axis=0)
        
        # 使用加权平均优化高斯卷积
        window = 21  # 增大窗口
        sigma = window/4.0  # 调整sigma
        x = np.linspace(-window//2, window//2, window)
        gaussian = np.exp(-(x**2)/(2*sigma**2))
        gaussian = gaussian/np.sum(gaussian)
        
        # 使用双向卷积
        smoothed = np.convolve(projection, gaussian, mode='same')
        smoothed = np.convolve(smoothed[::-1], gaussian, mode='same')[::-1]
        
        # 动态阈值
        gradient = np.gradient(smoothed)
        threshold = np.percentile(np.abs(gradient), 95)  # 使用95百分位数
        
        # 边缘位置平滑
        edge_positions = []
        for i in range(len(gradient)-1, window, -1):
            if abs(gradient[i]) > threshold:
                local_region = gradient[max(0, i-7):min(len(gradient), i+8)]
                local_offset = np.argmax(np.abs(local_region))
                edge_positions.append(i - 7 + local_offset)
                if len(edge_positions) >= 5:  # 取多个候选点
                    break
        
        if edge_positions:
            return int(np.median(edge_positions))  # 使用中位数
        return None
    
    def initialize_with_reference(self, reference_img):
        # 转换为灰度图
        gray = cv2.cvtColor(reference_img, cv2.COLOR_BGR2GRAY)
        
        # 增强对比度
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)
        
        # 检测垂直线
        binary = detect_vertical_lines(gray)
        
        # 找到右边缘
        right_edge = self.find_right_edge(binary)
        if right_edge is not None:
            self.reference_right_edge = right_edge
            # 设置裁剪宽度（从右边缘向左的距离）
            self.crop_width = right_edge - 100  # 可以根据需要调整这个值
    
    def crop_image(self, img):
        if self.crop_width is None:
            return img
            
        # 转换为灰度图并增强对比度
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)
        
        # 检测垂直线
        binary = detect_vertical_lines(gray)
        
        # 找到当前帧的右边缘
        current_right_edge = self.find_right_edge(binary)
        
        if current_right_edge is None:
            return img
            
        # 更新历史记录并获取平滑后的边缘位置
        self.update_edge_history(current_right_edge)
        smoothed_edge = self.get_smoothed_edge()
        
        # 动态调整裁剪区域
        start_x = max(0, smoothed_edge - self.crop_width)
        end_x = smoothed_edge
        
        # 添加边界检查
        if end_x - start_x < 100:  # 最小宽度
            start_x = max(0, end_x - 100)
            
        cropped = img[:, start_x:end_x]
        return cropped

def detect_vertical_lines(gray_img):
    # 使用Scharr算子替代Sobel，增强垂直边缘检测
    scharrx = cv2.Scharr(gray_img, cv2.CV_64F, 1, 0)
    scharrx = cv2.convertScaleAbs(scharrx)
    
    # 使用Otsu自适应阈值
    _, binary = cv2.threshold(scharrx, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    
    # 改进形态学操作
    kernel_vertical = cv2.getStructuringElement(cv2.MORPH_RECT, (21, 3))
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel_vertical)
    binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel_vertical)
    
    return binary

def process_single_image(args):
    img_path, cropper = args
    try:
        # 读取当前帧
        current_frame = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        
        if current_frame is None:
            return None
        
        # 裁剪图像
        cropped_frame = cropper.crop_image(current_frame)
        
        return cropped_frame
    except Exception as e:
        print(f"处理图片 {img_path} 时发生错误: {str(e)}")
        return None

def process_images_to_video(input_path, output_path, sample_count=4000):
    try:
        # 获取所有bmp文件
        image_files = [f for f in os.listdir(input_path) if f.lower().endswith('.bmp')]
        if not image_files:
            print(f"错误：在 {input_path} 中没有找到BMP文件")
            return
        image_files.sort()
        
        total_images = len(image_files)
        if total_images < sample_count:
            sample_count = total_images
        step = total_images // sample_count
        
        # 读取第一张图片作为参考帧
        first_image_path = os.path.join(input_path, image_files[0])
        print(f"正在读取第一张图片: {first_image_path}")
        reference_frame = cv2.imdecode(np.fromfile(first_image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        
        if reference_frame is None:
            print(f"错误：无法读取图片 {first_image_path}")
            return
        
        # 创建裁剪器并初始化
        cropper = ImageCropper()
        cropper.initialize_with_reference(reference_frame)
        
        # 获取裁剪后的尺寸
        first_processed = cropper.crop_image(reference_frame)
        height, width = first_processed.shape[:2]
        
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # 准备并行处理的参数
        selected_files = [os.path.join(input_path, image_files[i]) for i in range(0, total_images, step)]
        process_args = [(f, cropper) for f in selected_files]
        
        # 获取CPU核心数
        num_cores = multiprocessing.cpu_count()
        print(f"使用 {num_cores} 个CPU核心进行并行处理")
        
        # 创建进程池
        with Pool(num_cores) as pool:
            # 使用tqdm显示进度
            processed_frames = list(tqdm(
                pool.imap(process_single_image, process_args),
                total=len(process_args),
                desc="处理图片"
            ))
        
        # 过滤掉处理失败的帧
        processed_frames = [frame for frame in processed_frames if frame is not None]
        
        if not processed_frames:
            print("错误：没有成功处理任何图片")
            return
        
        # 创建视频写入器
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, 30.0, (width, height))
        
        # 写入处理后的帧
        for frame in processed_frames:
            out.write(frame)
        
        out.release()
        print("视频生成完成！")
        
    except Exception as e:
        print(f"发生错误: {str(e)}")

if __name__ == '__main__':
    # 使用原始字符串表示路径
    input_path = r"F:\李双成\B-1\200-4"  # 使用原始字符串
    output_path = r"F:\test\output_video_stabilized.mp4"  # 使用原始字符串
    process_images_to_video(input_path, output_path)
